# -*- coding: utf-8 -*-
# @Author: lidongdong
# @time  : 19-1-24 上午10:07
# @file  : data_loader.py

import numpy as np
from data_utils import *


class CocoLoader(object):
    def __init__(self, image_h5_filename, caption_json_filename, mode):
        self.mode = mode
        self.image_h5_filename = image_h5_filename
        self.caption_json_filename = caption_json_filename
        self.images = load_image(self.image_h5_filename)
        self.content = load_caption(self.caption_json_filename)     # content 下包含了 captions,caption_splits,image,index
        self.caption_index = 0
        self.caption_max = len(self.content)

    def reset(self):
        self.caption_index = 0

    def next(self, batch_size=32, index=0):
        """get next batch data

        # Returns:
            batched_images: 64, 224, 224, 3
            batched_captions:
            caption_lens:
        """
        end = self.caption_index + batch_size

        if end >= self.caption_max:
            self.reset()
            raise StopIteration

        # get images
        images = self.images[self.caption_index: end, :, :, :]
        batched_images = process_images(images, crop_scale_no=0)
        # get captions
        captions = self.content[self.caption_index: end]
        self.caption_index += batch_size

        temp = [caption["caption_splits"][index] for caption in captions]
        temp = [map(lambda x: x.encode("utf-8"), y) for y in temp]
        captions = temp
        max_len = max(map(lambda x: len(x), captions))
        caption_lens = []

        for i in range(len(captions)):
            caption_len = len(captions[i])
            # append <eos> for each caption
            caption_lens.append(caption_len + 1)
            captions[i].extend(["<eos>"] + ["<pad>"] * (max_len - caption_len))

        captions = map(lambda x: map(lambda z: z.lower(), x), captions)
        # append <eos> for each caption
        max_len += 1
        batched_images = np.asarray(batched_images)
        batched_captions = np.asarray(captions)
        caption_lens = np.asarray(caption_lens)
        return batched_images, batched_captions, caption_lens


if __name__ == '__main__':
    cocoloader = CocoLoader("../data/train.h5", "../data/train.json", "train")
    batched_images, batched_captions, caption_lens = cocoloader.next()

    print batched_images.shape
    print batched_captions.shape
    print batched_captions[0]
    print caption_lens
